#!/usr/bin/env python3
"""
Stable Video Diffusion 模型下载脚本 (镜像版本)
使用 HuggingFace 镜像站点下载模型
"""

import os
import sys
import yaml
import logging
from pathlib import Path

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class MirrorModelDownloader:
    def __init__(self, config_path='config/config.yaml'):
        """
        初始化模型下载器
        :param config_path: 配置文件路径
        """
        self.config = self._load_config(config_path)
        if not self.config:
            logger.error("配置文件加载失败，请先运行 setup_environment.py")
            return
            
        self.model_name = self.config['model']['name']
        self.cache_dir = self.config['model']['cache_dir']
        self.use_safetensors = self.config['model']['use_safetensors']

    def _load_config(self, config_path):
        """加载YAML配置文件"""
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                return yaml.safe_load(f)
        except FileNotFoundError:
            logger.error(f"配置文件未找到: {config_path}")
            return None
        except Exception as e:
            logger.error(f"加载配置文件时出错: {e}")
            return None

    def download_with_mirror(self):
        """
        使用镜像站点下载模型
        """
        if not self.config:
            return False

        logger.info("="*50)
        logger.info(f"使用镜像站点下载模型: {self.model_name}")
        logger.info(f"模型将保存到: {self.cache_dir}")
        logger.info("="*50)
        
        # 确保缓存目录存在
        Path(self.cache_dir).mkdir(exist_ok=True)

        # 设置镜像站点环境变量
        mirror_sites = [
            "https://hf-mirror.com",  # 常用的 HuggingFace 镜像
            "https://huggingface-mirror.com",
            "https://mirrors.huggingface.co"
        ]
        
        for mirror in mirror_sites:
            logger.info(f"\n尝试使用镜像站点: {mirror}")
            
            # 设置环境变量，让 huggingface_hub 使用镜像
            os.environ['HF_ENDPOINT'] = mirror
            
            try:
                from huggingface_hub import snapshot_download
                
                # 尝试通过镜像下载
                snapshot_download(
                    repo_id=self.model_name,
                    cache_dir=self.cache_dir,
                    local_dir_use_symlinks=False,
                    allow_patterns=["*.json", "*.safetensors" if self.use_safetensors else "*.bin"],
                    ignore_patterns=["*.ckpt", "*.pt", "*.png"],
                    token=os.environ.get("HUGGINGFACE_TOKEN"),
                    endpoint=mirror  # 显式指定端点
                )
                
                logger.info(f"\n✅ 通过镜像站点 {mirror} 下载成功！")
                logger.info(f"模型文件位于: {self.cache_dir}")
                
                # 清除环境变量
                if 'HF_ENDPOINT' in os.environ:
                    del os.environ['HF_ENDPOINT']
                
                return True
                
            except Exception as e:
                logger.warning(f"镜像站点 {mirror} 下载失败: {e}")
                continue
        
        # 如果所有镜像都失败，清除环境变量
        if 'HF_ENDPOINT' in os.environ:
            del os.environ['HF_ENDPOINT']
        
        logger.error("\n❌ 所有镜像站点都无法访问")
        logger.error("请尝试以下替代方案：")
        logger.error("1. 配置代理服务器")
        logger.error("2. 手动下载模型文件")
        logger.error("3. 使用提供的手动下载指南")
        
        return False

    def provide_manual_download_instructions(self):
        """提供手动下载指南"""
        print("\n" + "="*60)
        print("手动下载指南")
        print("="*60)
        print(f"\n模型名称: {self.model_name}")
        print(f"目标目录: {self.cache_dir}")
        print("\n步骤 1: 在可以访问 HuggingFace 的电脑上下载模型")
        print("-------------------------------------------------------")
        print("在您的本地电脑（能访问 HuggingFace 的）运行：")
        print(f"""
# 安装 huggingface-cli
pip install huggingface-hub

# 下载模型到本地
huggingface-cli download {self.model_name} \\
    --local-dir ./svd_model_files \\
    --local-dir-use-symlinks False
""")
        
        print("\n步骤 2: 打包模型文件")
        print("-------------------------------------------------------")
        print("""
# 在下载完成后，打包文件
tar -czf svd_model.tar.gz svd_model_files/
""")
        
        print("\n步骤 3: 上传到服务器")
        print("-------------------------------------------------------")
        print("""
# 使用 scp 或其他方式上传到服务器
scp svd_model.tar.gz user@your-server:/path/to/deployment/
""")
        
        print("\n步骤 4: 在服务器上解压")
        print("-------------------------------------------------------")
        print(f"""
# 在服务器上解压到正确的位置
cd /path/to/deployment
tar -xzf svd_model.tar.gz
mv svd_model_files/* {self.cache_dir}/
""")
        
        print("\n" + "="*60)

if __name__ == "__main__":
    downloader = MirrorModelDownloader()
    
    if downloader.config:
        # 先尝试镜像下载
        success = downloader.download_with_mirror()
        
        if not success:
            # 如果失败，提供手动下载指南
            downloader.provide_manual_download_instructions()
            sys.exit(1)
